{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-02-27T03:02:21.511850Z",
     "start_time": "2024-02-27T03:02:18.305513Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([1, 3, 32, 32])"
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "# 假设 x 是一个输入张量\n",
    "x = torch.randn(1, 3, 16, 16)  # 3D 输入数据，形状为 mini-batch x channels x height x width\n",
    "\n",
    "# 使用双线性插值法将 x 上采样到指定的尺寸\n",
    "upsampled_x = torch.nn.functional.interpolate(x, size=(32, 32), mode='bilinear', align_corners=False)\n",
    "upsampled_x.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "\n",
    "\n",
    "class Upsample(nn.Module):\n",
    "    def __init__(self, input_feat,out_feat):\n",
    "        super(Upsample, self).__init__()\n",
    "\n",
    "        self.body = nn.Sequential(  # nn.Conv2d(n_feat, n_feat*2, kernel_size=3, stride=1, padding=1, bias=False),\n",
    "            # dw\n",
    "            nn.Conv2d(input_feat, input_feat, kernel_size=3, stride=1, padding=1, groups=input_feat, bias=False, ),\n",
    "            # pw-linear\n",
    "            nn.Conv2d(input_feat, out_feat * 4, 1, 1, 0, bias=False),\n",
    "            # nn.BatchNorm2d(n_feat*2),\n",
    "            # nn.Hardswish(),\n",
    "            nn.PixelShuffle(2))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.body(x)\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-27T03:07:15.761196Z",
     "start_time": "2024-02-27T03:07:15.757370Z"
    }
   },
   "id": "867050a6cacd7aee",
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "y = torch.randn(1, 16, 8, 8)\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-27T03:08:48.681830Z",
     "start_time": "2024-02-27T03:08:48.679647Z"
    }
   },
   "id": "82e4f184b9df16b7",
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([1, 32, 16, 16])"
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "upsample = Upsample(16, 32)\n",
    "result = upsample(y)\n",
    "result.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-27T03:08:49.942280Z",
     "start_time": "2024-02-27T03:08:49.921168Z"
    }
   },
   "id": "a4597cbda39bf522",
   "execution_count": 8
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "class Downsample(nn.Module):\n",
    "    def __init__(self, input_feat,out_feat):\n",
    "        super(Downsample, self).__init__()\n",
    "\n",
    "        self.body = nn.Sequential(  # nn.Conv2d(n_feat, n_feat // 2, kernel_size=3, stride=1, padding=1, bias=False),\n",
    "            # dw\n",
    "            nn.Conv2d(input_feat, input_feat, kernel_size=3, stride=1, padding=1, groups=input_feat, bias=False, ),\n",
    "            # pw-linear\n",
    "            nn.Conv2d(input_feat, out_feat // 4, 1, 1, 0, bias=False),\n",
    "            # nn.BatchNorm2d(n_feat // 2),\n",
    "            # nn.Hardswish(),\n",
    "            nn.PixelUnshuffle(2))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.body(x)\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-27T03:09:39.599165Z",
     "start_time": "2024-02-27T03:09:39.596331Z"
    }
   },
   "id": "296103df389bb6ec",
   "execution_count": 9
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([1, 8, 4, 4])"
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "downsample = Downsample(16, 8)\n",
    "result_2 = downsample(y)\n",
    "result_2.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-27T03:10:23.957493Z",
     "start_time": "2024-02-27T03:10:23.943940Z"
    }
   },
   "id": "865b4a4d5f346dfd",
   "execution_count": 10
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "[0.0,\n 0.021739130839705467,\n 0.043478261679410934,\n 0.06521739065647125,\n 0.08695652335882187,\n 0.10869565606117249,\n 0.1304347813129425,\n 0.15217392146587372,\n 0.17391304671764374,\n 0.19565217196941376,\n 0.21739131212234497,\n 0.239130437374115,\n 0.260869562625885,\n 0.28260868787765503,\n 0.30434781312942505,\n 0.32608693838119507,\n 0.3478260636329651,\n 0.3695652186870575,\n 0.3913043439388275,\n 0.41304346919059753,\n 0.43478262424468994,\n 0.45652174949645996,\n 0.47826087474823,\n 0.5]"
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "depths=[2, 2, 18, 2]\n",
    "dpr = [\n",
    "    x.item() for x in torch.linspace(0, 0.5, sum(depths))\n",
    "] \n",
    "dpr"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T04:27:25.458820Z",
     "start_time": "2024-02-28T04:27:25.445448Z"
    }
   },
   "id": "89d2ae6180006396",
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "0"
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(depths[:0])"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T04:28:13.088757Z",
     "start_time": "2024-02-28T04:28:13.085098Z"
    }
   },
   "id": "74dcc431ce4de652",
   "execution_count": 12
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "2"
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(depths[:1])"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T04:28:24.383871Z",
     "start_time": "2024-02-28T04:28:24.380341Z"
    }
   },
   "id": "995615c65977d205",
   "execution_count": 13
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "[0.0, 0.021739130839705467]"
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dpr[0:2]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T04:28:33.832713Z",
     "start_time": "2024-02-28T04:28:33.829294Z"
    }
   },
   "id": "9390ca8a1cfa29fb",
   "execution_count": 14
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "class BasicConv2d(nn.Module):\n",
    "    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1):\n",
    "        super(BasicConv2d, self).__init__()\n",
    "\n",
    "        self.conv = nn.Conv2d(in_planes, out_planes,\n",
    "                              kernel_size=kernel_size, stride=stride,\n",
    "                              padding=padding, dilation=dilation, bias=False)\n",
    "        self.bn = nn.BatchNorm2d(out_planes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        x = self.bn(x)\n",
    "        return x"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T05:19:59.614655Z",
     "start_time": "2024-02-28T05:19:59.611398Z"
    }
   },
   "id": "410f1be10900d030",
   "execution_count": 15
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([1, 256, 64, 64])"
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(1, 1024, 64, 64)\n",
    "conv1 = BasicConv2d(1024, 256, 1, stride=1, padding=0, dilation=1)\n",
    "x = conv1(x)\n",
    "x.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T07:29:48.516723Z",
     "start_time": "2024-02-28T07:29:48.486442Z"
    }
   },
   "id": "bf6d54c0aec6fe5",
   "execution_count": 83
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([1, 256, 64, 64])"
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "\n",
    "y = up(x)\n",
    "y.shape\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T05:42:44.854322Z",
     "start_time": "2024-02-28T05:42:44.848439Z"
    }
   },
   "id": "5c2fae9428c4f331",
   "execution_count": 68
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "tensor([[[[ 0.0172,  0.0172,  0.0172,  ...,  0.0172,  0.0172,  0.0172],\n          [ 0.0172,  0.0172,  0.0172,  ...,  0.0172,  0.0172,  0.0172],\n          [ 0.0172,  0.0172,  0.0172,  ...,  0.0172,  0.0172,  0.0172],\n          ...,\n          [ 0.0172,  0.0172,  0.0172,  ...,  0.0172,  0.0172,  0.0172],\n          [ 0.0172,  0.0172,  0.0172,  ...,  0.0172,  0.0172,  0.0172],\n          [ 0.0172,  0.0172,  0.0172,  ...,  0.0172,  0.0172,  0.0172]],\n\n         [[-0.0098, -0.0098, -0.0098,  ..., -0.0098, -0.0098, -0.0098],\n          [-0.0098, -0.0098, -0.0098,  ..., -0.0098, -0.0098, -0.0098],\n          [-0.0098, -0.0098, -0.0098,  ..., -0.0098, -0.0098, -0.0098],\n          ...,\n          [-0.0098, -0.0098, -0.0098,  ..., -0.0098, -0.0098, -0.0098],\n          [-0.0098, -0.0098, -0.0098,  ..., -0.0098, -0.0098, -0.0098],\n          [-0.0098, -0.0098, -0.0098,  ..., -0.0098, -0.0098, -0.0098]],\n\n         [[ 0.0138,  0.0138,  0.0138,  ...,  0.0138,  0.0138,  0.0138],\n          [ 0.0138,  0.0138,  0.0138,  ...,  0.0138,  0.0138,  0.0138],\n          [ 0.0138,  0.0138,  0.0138,  ...,  0.0138,  0.0138,  0.0138],\n          ...,\n          [ 0.0138,  0.0138,  0.0138,  ...,  0.0138,  0.0138,  0.0138],\n          [ 0.0138,  0.0138,  0.0138,  ...,  0.0138,  0.0138,  0.0138],\n          [ 0.0138,  0.0138,  0.0138,  ...,  0.0138,  0.0138,  0.0138]],\n\n         ...,\n\n         [[ 0.0320,  0.0320,  0.0320,  ...,  0.0320,  0.0320,  0.0320],\n          [ 0.0320,  0.0320,  0.0320,  ...,  0.0320,  0.0320,  0.0320],\n          [ 0.0320,  0.0320,  0.0320,  ...,  0.0320,  0.0320,  0.0320],\n          ...,\n          [ 0.0320,  0.0320,  0.0320,  ...,  0.0320,  0.0320,  0.0320],\n          [ 0.0320,  0.0320,  0.0320,  ...,  0.0320,  0.0320,  0.0320],\n          [ 0.0320,  0.0320,  0.0320,  ...,  0.0320,  0.0320,  0.0320]],\n\n         [[ 0.0149,  0.0149,  0.0149,  ...,  0.0149,  0.0149,  0.0149],\n          [ 0.0149,  0.0149,  0.0149,  ...,  0.0149,  0.0149,  0.0149],\n          [ 0.0149,  0.0149,  0.0149,  ...,  0.0149,  0.0149,  0.0149],\n          ...,\n          [ 0.0149,  0.0149,  0.0149,  ...,  0.0149,  0.0149,  0.0149],\n          [ 0.0149,  0.0149,  0.0149,  ...,  0.0149,  0.0149,  0.0149],\n          [ 0.0149,  0.0149,  0.0149,  ...,  0.0149,  0.0149,  0.0149]],\n\n         [[ 0.0179,  0.0179,  0.0179,  ...,  0.0179,  0.0179,  0.0179],\n          [ 0.0179,  0.0179,  0.0179,  ...,  0.0179,  0.0179,  0.0179],\n          [ 0.0179,  0.0179,  0.0179,  ...,  0.0179,  0.0179,  0.0179],\n          ...,\n          [ 0.0179,  0.0179,  0.0179,  ...,  0.0179,  0.0179,  0.0179],\n          [ 0.0179,  0.0179,  0.0179,  ...,  0.0179,  0.0179,  0.0179],\n          [ 0.0179,  0.0179,  0.0179,  ...,  0.0179,  0.0179,  0.0179]]]],\n       grad_fn=<UpsampleBilinear2DBackward0>)"
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T05:36:15.166798Z",
     "start_time": "2024-02-28T05:36:15.162148Z"
    }
   },
   "id": "60f73b56abd29bfd",
   "execution_count": 60
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([1, 64, 64, 1024])"
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(1, 64, 64, 256)\n",
    "x_2 = torch.randn(1, 64, 64, 256)\n",
    "x_5 = torch.randn(1, 64, 64, 256)\n",
    "x_6 = torch.randn(1, 64, 64, 256)\n",
    "result = torch.cat([x, x_2, x_5, x_6], 3)\n",
    "result.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T07:24:56.908512Z",
     "start_time": "2024-02-28T07:24:56.888054Z"
    }
   },
   "id": "172df5ebcee7c327",
   "execution_count": 79
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from torchsummary import summary\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Enhancer(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(Enhancer, self).__init__()\n",
    "\n",
    "        self.relu = nn.LeakyReLU(0.2, inplace=True)\n",
    "\n",
    "        self.tanh = nn.Tanh()\n",
    "\n",
    "        self.refine1 = nn.Conv2d(in_channels, 20, kernel_size=3, stride=1, padding=1)\n",
    "        self.refine2 = nn.Conv2d(20, 20, kernel_size=3, stride=1, padding=1)\n",
    "\n",
    "        self.conv1010 = nn.Conv2d(20, 1, kernel_size=1, stride=1, padding=0)\n",
    "        self.conv1020 = nn.Conv2d(20, 1, kernel_size=1, stride=1, padding=0)\n",
    "        self.conv1030 = nn.Conv2d(20, 1, kernel_size=1, stride=1, padding=0)\n",
    "        self.conv1040 = nn.Conv2d(20, 1, kernel_size=1, stride=1, padding=0)\n",
    "\n",
    "        self.refine3 = nn.Conv2d(20 + 4, out_channels, kernel_size=3, stride=1, padding=1)\n",
    "        self.upsample = nn.functional.interpolate\n",
    "\n",
    "        self.batch1 = nn.InstanceNorm2d(100, affine=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        dehaze = self.relu((self.refine1(x)))\n",
    "        dehaze = self.relu((self.refine2(dehaze)))\n",
    "        shape_out = dehaze.data.size()\n",
    "\n",
    "        shape_out = shape_out[2:4]\n",
    "\n",
    "        x101 = F.avg_pool2d(dehaze, 32)\n",
    "\n",
    "        x102 = F.avg_pool2d(dehaze, 16)\n",
    "\n",
    "        x103 = F.avg_pool2d(dehaze, 8)\n",
    "\n",
    "        x104 = F.avg_pool2d(dehaze, 4)\n",
    "\n",
    "        x1010 = self.upsample(self.relu(self.conv1010(x101)), size=shape_out)\n",
    "        x1020 = self.upsample(self.relu(self.conv1020(x102)), size=shape_out)\n",
    "        x1030 = self.upsample(self.relu(self.conv1030(x103)), size=shape_out)\n",
    "        x1040 = self.upsample(self.relu(self.conv1040(x104)), size=shape_out)\n",
    "\n",
    "        dehaze = torch.cat((x1010, x1020, x1030, x1040, dehaze), 1)\n",
    "        dehaze = self.tanh(self.refine3(dehaze))\n",
    "\n",
    "        return dehaze\n",
    "    "
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:21:46.533713Z",
     "start_time": "2024-02-28T08:21:46.528051Z"
    }
   },
   "id": "86bc8d7f44d89c2d",
   "execution_count": 86
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "81701900ff85bfe9"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1          [1, 20, 256, 256]             560\n",
      "         LeakyReLU-2          [1, 20, 256, 256]               0\n",
      "            Conv2d-3          [1, 20, 256, 256]           3,620\n",
      "         LeakyReLU-4          [1, 20, 256, 256]               0\n",
      "            Conv2d-5               [1, 1, 8, 8]              21\n",
      "         LeakyReLU-6               [1, 1, 8, 8]               0\n",
      "            Conv2d-7             [1, 1, 16, 16]              21\n",
      "         LeakyReLU-8             [1, 1, 16, 16]               0\n",
      "            Conv2d-9             [1, 1, 32, 32]              21\n",
      "        LeakyReLU-10             [1, 1, 32, 32]               0\n",
      "           Conv2d-11             [1, 1, 64, 64]              21\n",
      "        LeakyReLU-12             [1, 1, 64, 64]               0\n",
      "           Conv2d-13           [1, 3, 256, 256]             651\n",
      "             Tanh-14           [1, 3, 256, 256]               0\n",
      "================================================================\n",
      "Total params: 4,915\n",
      "Trainable params: 4,915\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.75\n",
      "Forward/backward pass size (MB): 43.08\n",
      "Params size (MB): 0.02\n",
      "Estimated Total Size (MB): 43.85\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "enhancer = Enhancer(3, 3)\n",
    "summary(enhancer, input_size=(3, 256, 256), batch_size=1, device=\"cpu\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:22:52.782510Z",
     "start_time": "2024-02-28T08:22:52.757271Z"
    }
   },
   "id": "b642626f2ce9aa9c",
   "execution_count": 88
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "!set BASICSR_JIT=True\n",
    "from basicsr.archs.module import ResBlock, WarpBlock\n",
    "\n",
    "\n",
    "class RIDCPDecoder(nn.Module):\n",
    "    def __init__(self,\n",
    "                 in_channel,\n",
    "                 max_depth,\n",
    "                 input_res=256,\n",
    "                 channel_query_dict=None,\n",
    "                 norm_type='gn',\n",
    "                 act_type='leakyrelu',\n",
    "                 only_residual=False,\n",
    "                 use_warp=True\n",
    "                 ):\n",
    "        super().__init__()\n",
    "        self.only_residual = only_residual\n",
    "        self.use_warp = use_warp\n",
    "        self.upsampler = nn.ModuleList()\n",
    "        self.warp = nn.ModuleList()\n",
    "        res = input_res // (2 ** max_depth)\n",
    "        for i in range(max_depth):\n",
    "            in_channel, out_channel = channel_query_dict[res], channel_query_dict[res * 2]\n",
    "            self.upsampler.append(\n",
    "                nn.Sequential(\n",
    "                    nn.Upsample(scale_factor=2),\n",
    "                    nn.Conv2d(in_channel, out_channel, 3, stride=1, padding=1),\n",
    "                    ResBlock(out_channel, out_channel, norm_type, act_type),\n",
    "                    ResBlock(out_channel, out_channel, norm_type, act_type),\n",
    "                )\n",
    "            )\n",
    "            self.warp.append(WarpBlock(out_channel))\n",
    "            res = res * 2\n",
    "\n",
    "    def forward(self, x):\n",
    "        # in x (batch_size, 256, w, h)\n",
    "        # in code_decoder_output list length 2 shape (batch_size, 256, w, h)\n",
    "        # out channel 256 -> 128 -> 64\n",
    "        for idx, m in enumerate(self.upsampler):\n",
    "            with torch.backends.cudnn.flags(enabled=False):\n",
    "                    x = m(x)\n",
    "        return x\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:28:23.189850Z",
     "start_time": "2024-02-28T08:28:23.168937Z"
    }
   },
   "id": "ea3caba9644b56b7",
   "execution_count": 96
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "          Upsample-1         [1, 256, 256, 256]               0\n",
      "            Conv2d-2         [1, 128, 256, 256]         295,040\n",
      "         GroupNorm-3         [1, 128, 256, 256]             256\n",
      "         NormLayer-4         [1, 128, 256, 256]               0\n",
      "              SiLU-5         [1, 128, 256, 256]               0\n",
      "          ActLayer-6         [1, 128, 256, 256]               0\n",
      "            Conv2d-7         [1, 128, 256, 256]         147,584\n",
      "         GroupNorm-8         [1, 128, 256, 256]             256\n",
      "         NormLayer-9         [1, 128, 256, 256]               0\n",
      "             SiLU-10         [1, 128, 256, 256]               0\n",
      "         ActLayer-11         [1, 128, 256, 256]               0\n",
      "           Conv2d-12         [1, 128, 256, 256]         147,584\n",
      "         ResBlock-13         [1, 128, 256, 256]               0\n",
      "        GroupNorm-14         [1, 128, 256, 256]             256\n",
      "        NormLayer-15         [1, 128, 256, 256]               0\n",
      "             SiLU-16         [1, 128, 256, 256]               0\n",
      "         ActLayer-17         [1, 128, 256, 256]               0\n",
      "           Conv2d-18         [1, 128, 256, 256]         147,584\n",
      "        GroupNorm-19         [1, 128, 256, 256]             256\n",
      "        NormLayer-20         [1, 128, 256, 256]               0\n",
      "             SiLU-21         [1, 128, 256, 256]               0\n",
      "         ActLayer-22         [1, 128, 256, 256]               0\n",
      "           Conv2d-23         [1, 128, 256, 256]         147,584\n",
      "         ResBlock-24         [1, 128, 256, 256]               0\n",
      "         Upsample-25         [1, 128, 512, 512]               0\n",
      "           Conv2d-26          [1, 64, 512, 512]          73,792\n",
      "        GroupNorm-27          [1, 64, 512, 512]             128\n",
      "        NormLayer-28          [1, 64, 512, 512]               0\n",
      "             SiLU-29          [1, 64, 512, 512]               0\n",
      "         ActLayer-30          [1, 64, 512, 512]               0\n",
      "           Conv2d-31          [1, 64, 512, 512]          36,928\n",
      "        GroupNorm-32          [1, 64, 512, 512]             128\n",
      "        NormLayer-33          [1, 64, 512, 512]               0\n",
      "             SiLU-34          [1, 64, 512, 512]               0\n",
      "         ActLayer-35          [1, 64, 512, 512]               0\n",
      "           Conv2d-36          [1, 64, 512, 512]          36,928\n",
      "         ResBlock-37          [1, 64, 512, 512]               0\n",
      "        GroupNorm-38          [1, 64, 512, 512]             128\n",
      "        NormLayer-39          [1, 64, 512, 512]               0\n",
      "             SiLU-40          [1, 64, 512, 512]               0\n",
      "         ActLayer-41          [1, 64, 512, 512]               0\n",
      "           Conv2d-42          [1, 64, 512, 512]          36,928\n",
      "        GroupNorm-43          [1, 64, 512, 512]             128\n",
      "        NormLayer-44          [1, 64, 512, 512]               0\n",
      "             SiLU-45          [1, 64, 512, 512]               0\n",
      "         ActLayer-46          [1, 64, 512, 512]               0\n",
      "           Conv2d-47          [1, 64, 512, 512]          36,928\n",
      "         ResBlock-48          [1, 64, 512, 512]               0\n",
      "================================================================\n",
      "Total params: 1,108,416\n",
      "Trainable params: 1,108,416\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 16.00\n",
      "Forward/backward pass size (MB): 4800.00\n",
      "Params size (MB): 4.23\n",
      "Estimated Total Size (MB): 4820.23\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "ridcp = RIDCPDecoder(\n",
    "    3,\n",
    "    2,\n",
    "    256,\n",
    "    {\n",
    "        8: 256,\n",
    "        16: 256,\n",
    "        32: 256,\n",
    "        64: 256,\n",
    "        128: 128,\n",
    "        256: 64,\n",
    "        512: 32,\n",
    "    },\n",
    "    'gn', 'silu', False, use_warp=True\n",
    ").cuda()\n",
    "\n",
    "summary(ridcp, input_size=(256, 128, 128), batch_size=1, device=\"cuda\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:54:14.226260Z",
     "start_time": "2024-02-28T08:54:13.877527Z"
    }
   },
   "id": "60132b272293f688",
   "execution_count": 115
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "class PALayer(nn.Module):\n",
    "    def __init__(self, channel):\n",
    "        super(PALayer, self).__init__()\n",
    "        self.pa = nn.Sequential(\n",
    "            nn.Conv2d(channel, channel // 8, 1, padding=0, bias=True),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(channel // 8, 1, 1, padding=0, bias=True),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        y = self.pa(x)\n",
    "        return x * y\n",
    "\n",
    "\n",
    "class CALayer(nn.Module):\n",
    "    def __init__(self, channel):\n",
    "        super(CALayer, self).__init__()\n",
    "        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n",
    "        self.ca = nn.Sequential(\n",
    "            nn.Conv2d(channel, channel // 8, 1, padding=0, bias=True),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(channel // 8, channel, 1, padding=0, bias=True),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        y = self.avg_pool(x)\n",
    "        y = self.ca(y)\n",
    "        return x * y\n",
    "\n",
    "\n",
    "class DehazeBlock(nn.Module):\n",
    "    def __init__(self, conv, dim, kernel_size):\n",
    "        super(DehazeBlock, self).__init__()\n",
    "\n",
    "        self.conv1 = conv(dim, dim, kernel_size, bias=True)\n",
    "        self.act1 = nn.ReLU(inplace=True)\n",
    "        self.conv2 = conv(dim, dim, kernel_size, bias=True)\n",
    "        self.calayer = CALayer(dim)\n",
    "        self.palayer = PALayer(dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        res = self.act1(self.conv1(x))\n",
    "        res = res + x\n",
    "        res = self.conv2(res)\n",
    "        res = self.calayer(res)\n",
    "        res = self.palayer(res)\n",
    "        res += x\n",
    "\n",
    "        return res\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:58:47.744926Z",
     "start_time": "2024-02-28T08:58:47.738187Z"
    }
   },
   "id": "1adc05d23e9676e6",
   "execution_count": 122
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1           [1, 256, 64, 64]          65,792\n",
      "              ReLU-2           [1, 256, 64, 64]               0\n",
      "            Conv2d-3           [1, 256, 64, 64]          65,792\n",
      " AdaptiveAvgPool2d-4             [1, 256, 1, 1]               0\n",
      "            Conv2d-5              [1, 32, 1, 1]           8,224\n",
      "              ReLU-6              [1, 32, 1, 1]               0\n",
      "            Conv2d-7             [1, 256, 1, 1]           8,448\n",
      "           Sigmoid-8             [1, 256, 1, 1]               0\n",
      "           CALayer-9           [1, 256, 64, 64]               0\n",
      "           Conv2d-10            [1, 32, 64, 64]           8,224\n",
      "             ReLU-11            [1, 32, 64, 64]               0\n",
      "           Conv2d-12             [1, 1, 64, 64]              33\n",
      "          Sigmoid-13             [1, 1, 64, 64]               0\n",
      "          PALayer-14           [1, 256, 64, 64]               0\n",
      "================================================================\n",
      "Total params: 156,513\n",
      "Trainable params: 156,513\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 4.00\n",
      "Forward/backward pass size (MB): 42.07\n",
      "Params size (MB): 0.60\n",
      "Estimated Total Size (MB): 46.67\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "dehaze_block = DehazeBlock(nn.Conv2d, 256, 1)\n",
    "dehaze_block = dehaze_block.cuda()\n",
    "summary(dehaze_block, input_size=(256, 64, 64), batch_size=1, device=\"cuda\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:58:48.433775Z",
     "start_time": "2024-02-28T08:58:48.413910Z"
    }
   },
   "id": "9afa20af09fa7762",
   "execution_count": 123
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "64"
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "channel_query_dict = {\n",
    "    8: 256,\n",
    "    16: 256,\n",
    "    32: 256,\n",
    "    64: 256,\n",
    "    128: 128,\n",
    "    256: 64,\n",
    "    512: 32,\n",
    "}\n",
    "res = 256 // (2 ** 2)\n",
    "res"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:55:54.600058Z",
     "start_time": "2024-02-28T08:55:54.597074Z"
    }
   },
   "id": "ff170671cd481a28",
   "execution_count": 116
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "in_channel, out_channel = channel_query_dict[res], channel_query_dict[res * 2]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:55:57.766681Z",
     "start_time": "2024-02-28T08:55:57.764339Z"
    }
   },
   "id": "44f283b45c4e9548",
   "execution_count": 117
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "(256, 128)"
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "in_channel, out_channel"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:56:06.456283Z",
     "start_time": "2024-02-28T08:56:06.452614Z"
    }
   },
   "id": "6ba3cbcfdc1509c6",
   "execution_count": 118
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "res = 128\n",
    "in_channel, out_channel = channel_query_dict[res], channel_query_dict[res * 2]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:56:49.123838Z",
     "start_time": "2024-02-28T08:56:49.121289Z"
    }
   },
   "id": "e79b3c3a6df73ca6",
   "execution_count": 119
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "(128, 64)"
     },
     "execution_count": 120,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "in_channel, out_channel"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T08:56:52.363534Z",
     "start_time": "2024-02-28T08:56:52.360959Z"
    }
   },
   "id": "fbb2920164a42bdd",
   "execution_count": 120
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "ecd920373c549a0b"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
